import json
import os

# 获取数据文件所在目录的绝对路径
cur = os.path.join(os.path.dirname(os.path.abspath(__file__)), '../')

class TransferData():
    def __init__(self):
        self.label_dict = json.load(open(cur + '/data/labels.json', 'r', encoding='utf-8'))
        self.seq_tag_dict = json.load(open(cur + '/data/tag2id.json', 'r', encoding='utf-8'))
        self.origin_path = os.path.join(cur, 'data_origin')
        self.train_filepath = os.path.join(cur, 'data/train.txt')

    def transfer(self):
        with open(file=self.train_filepath,
                  mode='w',
                  encoding='utf-8') as fw:
            for root, dirs, files in os.walk(self.origin_path):
                for file in files:
                    filepath = os.path.join(root, file)
                    if 'original' in filepath:
                        label_filepath = filepath.replace('.txtoriginal', '')
                        res_dict = self.read_label_text(label_filepath)
                        with open(filepath, 'r', encoding='utf-8') as f:
                            content = f.read().strip()
                            if content and content[-1] not in ['。', '?', '!', '！', '？']:
                                content += '。'
                        for indx, char in enumerate(content):
                            char_label = res_dict.get(indx, 'O')
                            fw.write(char + '\t' + char_label + '\n')

    def read_label_text(self, label_filepath):
        res_dict = {}
        for line in open(label_filepath, 'r', encoding='utf-8'):
            res = line.strip().split('\t')
            if res:
                start = int(res[1])
                end = int(res[2])
                label = res[3]
                label_tag = self.label_dict.get(label)
                for i in range(start, end + 1):
                    if i == start:
                        tag = 'B-' + label_tag
                    else:
                        tag = 'I-' + label_tag
                    res_dict[i] = tag
        return res_dict

if __name__ == '__main__':
    handler = TransferData()
    handler.transfer()
